"""
Python代码执行组件

提供Python代码的安全执行和语法验证功能
"""

import asyncio
import io
import traceback
import multiprocessing
from concurrent.futures import ProcessPoolExecutor
from typing import Any, List, Optional, Dict
from contextlib import redirect_stdout, redirect_stderr

import mcp.types as types

from .base import BaseMCPComponent


class PythonExecutorComponent(BaseMCPComponent):
    """Python代码执行组件"""
    
    def __init__(self, max_workers: Optional[int] = None):
        """初始化Python执行组件"""
        self._max_workers = max_workers or min(4, multiprocessing.cpu_count())
        self._executor: Optional[ProcessPoolExecutor] = None
        self._tools = {}
        super().__init__(name="python_executor")
    
    def initialize(self) -> None:
        """初始化组件（预加载进程池）"""
        if self._executor is None:
            self._executor = ProcessPoolExecutor(max_workers=self._max_workers)
    
    def is_available(self) -> bool:
        """检查组件是否可用"""
        return True
    
    @staticmethod
    def _execute_code(code: str, timeout: float = 30) -> Dict[str, Any]:
        """安全执行Python代码（静态方法，避免序列化问题）"""
        stdout_capture = io.StringIO()
        stderr_capture = io.StringIO()
        
        try:
            exec_globals = {"__builtins__": __builtins__, "print": print}
            exec_locals = {}
            
            with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture):
                exec(compile(code, "<string>", "exec"), exec_globals, exec_locals)
            
            output = stdout_capture.getvalue()
            stderr_out = stderr_capture.getvalue()
            if stderr_out:
                output += f"\n[stderr]: {stderr_out}"
            
            return {
                "success": True,
                "output": output,
                "return_value": str(exec_locals) if exec_locals else None
            }
        except Exception as e:
            return {
                "success": False,
                "output": stdout_capture.getvalue(),
                "error": f"{type(e).__name__}: {str(e)}\n{traceback.format_exc()}"
            }
    
    @staticmethod
    def _validate_syntax(code: str) -> Dict[str, Any]:
        """验证Python代码语法（静态方法）"""
        try:
            compile(code, "<string>", "exec")
            return {"valid": True, "message": "代码语法正确"}
        except SyntaxError as e:
            return {"valid": False, "error": f"语法错误在第 {e.lineno} 行: {e.msg}"}
        except Exception as e:
            return {"valid": False, "error": f"验证错误: {str(e)}"}
    
    def _get_tools_definition(self) -> List[types.Tool]:
        """获取工具定义"""
        return [
            types.Tool(
                name="execute_python",
                description="执行Python代码并返回结果",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "code": {"type": "string", "description": "要执行的Python代码"},
                        "timeout": {"type": "number", "description": "超时时间（秒）", "default": 30}
                    },
                    "required": ["code"]
                }
            ),
            types.Tool(
                name="validate_python",
                description="验证Python代码语法",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "code": {"type": "string", "description": "要验证的Python代码"}
                    },
                    "required": ["code"]
                }
            )
        ]
    
    async def handle_tool_call(self, name: str, arguments: Any) -> List[types.TextContent]:
        """处理工具调用"""
        if name == "execute_python":
            code = arguments.get("code", "")
            if not code:
                return [types.TextContent(type="text", text="✗ 错误: 未提供代码")]
            
            timeout = min(arguments.get("timeout", 30), 300)
            
            # 异步执行
            loop = asyncio.get_event_loop()
            if self._executor is None:
                self.initialize()
            
            result = await loop.run_in_executor(
                self._executor,
                PythonExecutorComponent._execute_code,
                code,
                timeout
            )
            
            if result["success"]:
                parts = []
                if result["output"]:
                    parts.append(f"输出:\n{result['output']}")
                if result["return_value"]:
                    parts.append(f"返回值: {result['return_value']}")
                
                text = "✓ 执行成功\n\n" + ("\n\n".join(parts) if parts else "代码执行成功，无输出")
            else:
                parts = []
                if result["output"]:
                    parts.append(f"输出:\n{result['output']}")
                if result["error"]:
                    parts.append(f"错误:\n{result['error']}")
                text = "✗ 执行失败\n\n" + "\n\n".join(parts)
            
            return [types.TextContent(type="text", text=text)]
        
        elif name == "validate_python":
            code = arguments.get("code", "")
            if not code:
                return [types.TextContent(type="text", text="✗ 错误: 未提供代码")]
            
            result = PythonExecutorComponent._validate_syntax(code)
            
            if result["valid"]:
                text = f"✓ {result['message']}"
            else:
                text = f"✗ 语法验证失败\n\n错误: {result['error']}"
            
            return [types.TextContent(type="text", text=text)]
        
        else:
            return [types.TextContent(type="text", text=f"✗ 未知工具: {name}")]
    
    def cleanup(self) -> None:
        """清理资源"""
        if self._executor:
            self._executor.shutdown(wait=True)
            self._executor = None
        super().cleanup()
